# -*- coding: utf-8 -*-
"""
Created on Tue Dec 01 21:05:42 2015

@author: Tim
"""
import matplotlib.pyplot as P
from pylab import *                ## import scientific database
close("all")                       ## close all windows
from netCDF4 import Dataset
import numpy as np
from mpl_toolkits.basemap import Basemap
from matplotlib.font_manager import FontProperties
from mpl_toolkits.axes_grid.inset_locator import inset_axes
from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes
from mpl_toolkits.axes_grid1.inset_locator import mark_inset
from wrf import getvar, interplevel, to_np, get_basemap, latlon_coords, CoordPair, vertcross
import glob
import cmocean
import cmocean.cm as cmo


import cartopy.crs as ccrs
import cartopy.feature
from cartopy.io.img_tiles import MapQuestOpenAerial
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D as Line
from matplotlib.patheffects import Stroke
import numpy as np
import shapely.geometry as sgeom
from shapely.ops import transform as geom_transform


g = 9.81

ll_lona = np.zeros(2)
ll_lata = np.zeros(2)
ul_lona = np.zeros(2)
ul_lata = np.zeros(2)
ur_lona = np.zeros(2)
ur_lata = np.zeros(2)
lr_lona = np.zeros(2)
lr_lata = np.zeros(2)

for i in np.arange(2):
    if i == 0:
        wrf_files = '/glade/scratch/tjuliano/fire/sagehen/WRF-Fire-LEAPHI/WRF/test/marshall/2021123012_maria/wrfinput_d01'
    elif i == 1:
        wrf_files = '/glade/scratch/tjuliano/fire/sagehen/WRF-Fire-LEAPHI/WRF/test/marshall/2021123012_maria/wrfinput_d02'

    print ('Reading WRF files...')

    s1 = Dataset(wrf_files)

    lat2 = getvar(s1,"lat",meta=False)
    lon2 = getvar(s1,"lon",meta=False)
    hgt = getvar(s1,"HGT",meta=False)
    hgt[hgt==0.0] = -100.
    if i == 0:
    	lat_d01 = lat2
    	lon_d01 = lon2
    	hgt_d01 = hgt
    elif i == 1:
    	lat_d02 = lat2
    	lon_d02 = lon2
    	hgt_d02 = hgt

    if i < 2:
	    ll_lon = lon2[0,0]
	    ll_lat = lat2[0,0]
	    ul_lon = lon2[-1,0]
	    ul_lat = lat2[-1,0]
	    ur_lon = lon2[-1,-1]
	    ur_lat = lat2[-1,-1]
	    lr_lon = lon2[0,-1]
	    lr_lat = lat2[0,-1]

	    ll_lona[i] = ll_lon
	    ll_lata[i] = ll_lat
	    ul_lona[i] = ul_lon
	    ul_lata[i] = ul_lat
	    ur_lona[i] = ur_lon
	    ur_lata[i] = ur_lat
	    lr_lona[i] = lr_lon
	    lr_lata[i] = lr_lat

#ll_lon = ll_lona[0]
#ll_lat = ll_lata[0]
#ur_lon = ur_lona[0]
#ur_lat = ur_lata[0]

ll_lon = -105.7
ll_lat = 39.55
ur_lon = -104.6
ur_lat = 40.3

# center point of overall domain
ref_lat = (ll_lat + ur_lat)/2.
ref_lon = -(abs(ll_lon) + abs(ur_lon))/2.
#ref_lon = 8.5

# create map subplots
# 1000 MB
fig = plt.figure(figsize=(20,15))
ax = fig.add_axes([0.05,0.05,0.9,0.9])
m = Basemap(llcrnrlon=ll_lon,llcrnrlat=ll_lat,urcrnrlon=ur_lon,urcrnrlat=ur_lat,lon_0=ref_lon,lat_0=ref_lat,projection='lcc',resolution='h')

x_d01, y_d01 = m(lon_d01,lat_d01)
x_d02, y_d02 = m(lon_d02,lat_d02)

my_cmap = cmo.deep
#my_cmap.set_under('lightskyblue')
#my_cmap2 = cmo.thermal

levels = np.arange(1200,3800,200)
CS2 = m.contourf(x_d01,y_d01,hgt_d01,levels,
                   cmap=my_cmap,
		   extend='max',
                   alpha=1.0,
                   origin='lower',
                   zorder=1)

CS2 = m.contourf(x_d02,y_d02,hgt_d02,levels,
                   cmap=my_cmap,
                   extend='max',
                   alpha=1.0,
                   origin='lower',
                   zorder=2)

#import geopandas as gpd
#data = gpd.read_file('tl_2021_08_prisecroads.shp')
#sorted = data[(data.RTTYP == 'I') | (data.RTTYP == 'U') | (data.FULLNAME == 'State Rte 101') | (data.FULLNAME == 'Rte 101')]
#sorted = data[data.FULLNAME == 'Hwy 36']
#print (sorted)
#m.add_geometries(sorted.geometry, linewidth=2.0, facecolor='none', edgecolor='gray',zorder=3)

m.readshapefile('tl_2021_08_prisecroads', 'roads', drawbounds = False)

for info, shape in zip(m.roads_info, m.roads):
#    print (info)
    if info['FULLNAME'] == 'US Hwy 36' or info['FULLNAME'] == 'I- 25' or info['FULLNAME'] == 'I- 70':
        x, y = zip(*shape) 
        m.plot(x, y, marker=None,color='m',lw=3.,zorder=3)


# customize map
#m.drawmapboundary(fill_color='aqua')
#m.drawcoastlines(linewidth=2.0)
m.drawparallels(np.arange(30.,94.,0.2),labels=[True,False,False,False],fontsize=28,linewidth=3.0)
m.drawmeridians(np.arange(-130.,-100.,0.2),labels=[False,False,False,True],fontsize=28,linewidth=3.0)
#m.drawstates(linewidth=2.0)
#m.drawcountries(linewidth=2.0)
#m.drawrivers(linewidth=3.0,color='blue')
#m.fillcontinents(color='coral',lake_color='aqua')

# get x/y coordinates from lat/lon
ll_x, ll_y = m(ll_lona, ll_lata)
ul_x, ul_y = m(ul_lona, ul_lata)
ur_x, ur_y = m(ur_lona, ur_lata)
lr_x, lr_y = m(lr_lona, lr_lata)

# left
m.plot([ll_x[1],ul_x[1]],[ll_y[1],ul_y[1]],linewidth=6,color='red',zorder=4)
# top
m.plot([ul_x[1],ur_x[1]],[ul_y[1],ur_y[1]],linewidth=6,color='red',zorder=4)
# right
m.plot([ur_x[1],lr_x[1]],[ur_y[1],lr_y[1]],linewidth=6,color='red',zorder=4)
# bottom
m.plot([lr_x[1],ll_x[1]],[lr_y[1],ll_y[1]],linewidth=6,color='red',zorder=4)
font0 = FontProperties()
font = font0.copy()
font.set_weight('semibold')
x1, y1 = m(-105.43,40.125)
plt.text(x1,y1,'d02\ndx=111m',fontsize=28,color='black',bbox=dict(facecolor='darkgray', edgecolor='white', alpha=1.0, boxstyle='round'),zorder=5)

x1, y1 = m(-105.17,39.87)
plt.text(x1,y1,'Hwy\n 36',fontsize=28,color='k',zorder=5)
x1, y1 = m(-104.96,40.1)
plt.text(x1,y1,'I-25',fontsize=28,color='k',zorder=5)
x1, y1 = m(-105.3,39.725)
plt.text(x1,y1,'I-70',fontsize=28,color='k',zorder=5)

##x1, y1 = m(-108.2,42.275)
##plt.text(x1,y1,'dx=1000m',fontsize=28,color='white')
#x2, y2 = m(28.0,74.0)
#x2, y2 = m(-4.0,68.1)
#plt.text(x2,y2,'d02\ndx=1km',fontsize=28,color='black',bbox=dict(facecolor='darkgray', edgecolor='white', alpha=1.0, boxstyle='round'),zorder=5)
##x1, y1 = m(-106.4,40.475)
##plt.text(x1,y1,'dx=111m',fontsize=28,color='white')

#x3, y3 = m(-121.4,46.3)
#plt.text(x3,y3,'Mt. Adams',fontsize=28,color='white')
#x4, y4 = m(-122.1,45.3)
#plt.text(x4,y4,'Mt. Hood',fontsize=28,color='white')
#x4, y4 = m(-121.7,44.7)
#plt.text(x4,y4,'Mt. Jefferson',fontsize=28,color='white')

#x4, y4 = m(-121.95,45.75)
#plt.text(x4,y4,'Columbia',fontsize=28,color='white')
#x4, y4 = m(-121.85,45.6)
#plt.text(x4,y4,'River',fontsize=28,color='white')

#x5, y5 = m(-120.8,46.1)
#plt.text(x5,y5,'Washington',fontsize=48,color='white')
#x5, y5 = m(-120.4,44.8)
#plt.text(x5,y5,'Oregon',fontsize=48,color='white')

### Locations
### ['ATC10', 'AV313', 'BLD01', 'CO109', 'E6155', 'E9298', 'E9688', 'G0194']
wrf_stations = ['ATC', 'AV3', 'BLD', 'CO1', 'E61', 'E92', 'E96', 'G01']
raws_lats = [39.98521,39.99083,39.97400,39.86567,39.95550,39.98950,39.86667,40.00167]
raws_lons = [-105.15463,-105.11033,-105.24600,-105.24060,-105.16800,-105.08883,-105.22550,-105.22100]
c = ['r','r','orange','orange','b','b','lightskyblue','lightskyblue']
mark = ['o','^','o','^','o','^','o','^']

for i in np.arange(len(raws_lats)):
    x1, y1 = m(raws_lons[i],raws_lats[i])
    plt.scatter(x1,y1,marker=mark[i],s=800,fc=c[i],ec='gray',linewidths=2.,zorder=4,label=wrf_stations[i])

x1, y1 = m(-105.230189,39.956029)
plt.scatter(x1,y1,s=1600,marker='*',fc='magenta',ec='gray',linewidths=2.,zorder=4,label='Ign')

plt.legend(loc='center right',bbox_to_anchor=(1.08,0.4),ncol=1,handlelength=2,fontsize=28,labelspacing=1.1)

#cbarax = fig.add_axes([0.8,0.05,0.04,0.9])
cbar = plt.colorbar(CS2, orientation='horizontal', fraction=0.045, pad=0.1)
cbar.set_label('Elevation (m)',size=36,labelpad=30)
cbar.ax.tick_params(labelsize=28)
#cbar.set_ticks([1500,2000,2500,3000,3500,4000])

# Define the coordinate system of the data we have from Natural Earth and
# acquire the 1:10m physical coastline shapefile.
geodetic = ccrs.Geodetic(globe=ccrs.Globe(datum='WGS84'))

# Create an inset GeoAxes showing the location of the Solomon Islands.
ll_lon = ll_lona[0]
ll_lat = ll_lata[0]
ur_lon = ur_lona[0]
ur_lat = ur_lata[0]
sub_ax = plt.axes([0.7, 0.8, 0.2, 0.2], projection=ccrs.LambertConformal(central_longitude=ref_lon, central_latitude=ref_lat))
sub_ax.set_extent([ll_lon-4, ur_lon+4, ll_lat-4, ur_lat+4])

# Make a nice border around the inset axes.
effect = Stroke(linewidth=4, foreground='wheat')
sub_ax.outline_patch.set_path_effects([effect])

# Add the land, coastlines and the extent of the Solomon Islands.
sub_ax.add_feature(cartopy.feature.STATES)
extent = (ll_lon, ur_lon, ll_lat, ur_lat)
extent_box = sgeom.box(extent[0], extent[2], extent[1], extent[3])
sub_ax.add_geometries([extent_box], ccrs.PlateCarree(), edgecolor='blue',
                      facecolor='none', linewidth=4, linestyle='--', zorder=3)

ll_lon = ll_lona[1]
ll_lat = ll_lata[1]
ur_lon = ur_lona[1]
ur_lat = ur_lata[1]
extent = (ll_lon, ur_lon, ll_lat, ur_lat)
extent_box = sgeom.box(extent[0], extent[2], extent[1], extent[3])
sub_ax.add_geometries([extent_box], ccrs.PlateCarree(), edgecolor='red',
                      facecolor='none', linewidth=2, linestyle='-', zorder=3)

sub_ax.text(-108.3,42.6,'d01',fontsize=18,color='b',zorder=5,transform=ccrs.PlateCarree())

#plt.tight_layout()
#plt.subplots_adjust(right=0.775)
plt.savefig('domains_d02.png',dpi=250,bbox_inches='tight')
